#!/usr/bin/env python

from PyQt4.QtGui import *
from PyQt4.QtCore import *
import FrameFirewallBase
import FirewallRuleEditBase


class FrameFirewall(FrameFirewallBase.Ui_FrameFirewall, QFrame):
    def __init__(self, connection, parent = None):
        super(FrameFirewall, self).__init__(parent)
        self.setupUi(self)

        self.connection = connection
        self.Firewall = self.connection.root.Modules.Firewall.Firewall()

        self.holdChainsSelection = False
        self.holdRulesSelection = False
        self.selectedAChain = False
        self.selectedARule = False
        self.selectFirstRule = True

    def getChains(self, tableName):
        '''
        getChains(tableName)
            Add all the chains on the selected table to the chains list. 
        '''
        self.chainsList = self.Firewall.listChains(tableName)
        self.lstChains.clear()

        for chain in self.chainsList :
            listItem = QListWidgetItem(chain[0] + ' (' + chain[1] + ')')
            self.lstChains.addItem(listItem)

    def formatRules(self, rawRulesList):
        '''
        formatRules(rawRulesList)
            Detects if the rule has no action and add "No Action" to it. 
        '''
        rulesList = []
        for rawRule in rawRulesList :
            if len(rawRule) == 4 :
                rawRule.insert(0, 'No action')
            del(rawRule[2])
            del(rawRule[4:])
            rulesList.append(rawRule)

        return rulesList

    def refreshChains(self):
        currentRow = self.lstChains.currentRow()
        if currentRow == self.lstChains.count() - 1 :
            currentRow = currentRow - 1

        self.on_btnListChains_clicked()
        self.lstChains.setCurrentRow(currentRow)

    def refreshRules(self, action):
        currentIndex = self.lstRules.currentRow()

        if action == 'del' :
            if currentIndex == self.lstRules.count() - 1 :
                currentIndex = currentIndex - 1

        elif action == 'up' :
            currentIndex = currentIndex - 1

        elif action == 'down' :
            currentIndex = currentIndex + 1

        elif action == 'add' :
            currentIndex = self.lstRules.count()

        elif action == 'change' :
            pass

        self.selectFirstRule = False
        self.on_lstChains_itemSelectionChanged()
        self.lstRules.setCurrentRow(currentIndex)
        self.selectFirstRule = True


    def moveRule(self, direction):

        tableName = str(self.cmbTables.currentText())
        currentChain = self.chainsList[self.lstChains.currentRow()][0]
        ruleID = self.lstRules.currentRow() + 1
        diffrence = 0

        if direction == 'up' :
            if ruleID == 1 :
                return
            diffrence = -1

        elif direction == 'down' :
            if ruleID == self.lstRules.count():
                return
            diffrence = 1

        else :
            raise TypeError

        newID = ruleID + diffrence

        if self.Firewall.moveRule(tableName, currentChain, str(ruleID), str(newID)) :
            self.refreshRules(direction)





    #####################################################################
    #                                                                   #
    #                      Handling signals                             #
    #                           Chains                                  #
    #####################################################################

    @pyqtSlot()
    def on_btnListChains_clicked(self):
        tableName = str(self.cmbTables.currentText())
        self.holdChainsSelection = True
        self.getChains(tableName)
        self.holdChainsSelection = False
        self.lstChains.setCurrentItem(self.lstChains.item(0))

    @pyqtSlot('QString')
    def on_cmbTables_currentIndexChanged(self, selection):
        tableName = str(selection)
        self.holdChainsSelection = True
        self.getChains(tableName)
        self.holdChainsSelection = False
        self.lstChains.setCurrentItem(self.lstChains.item(0))

    def on_lstChains_itemSelectionChanged(self):
        if self.holdChainsSelection == True :
            return

        if self.selectedAChain == False :
            self.cmbChainsAction.removeItem(0)
            self.selectedAChain = True

        currentTable = str(self.cmbTables.currentText())
        currentChain = self.chainsList[self.lstChains.currentRow()][0]
        currentPolicy = self.chainsList[self.lstChains.currentRow()][1]

        actionsList = ['ACCEPT', 'DROP', 'User defined']

        if currentPolicy == 'User defined' :
            self.cmbChainsAction.setEnabled(False)
            self.btnChainsApply.setEnabled(False)
            self.lstRules.clear()

        self.cmbChainsAction.setEnabled(True)
        self.btnChainsApply.setEnabled(True)

        action = actionsList.index(currentPolicy)
        self.cmbChainsAction.setCurrentIndex(action)

        rawRules = self.Firewall.listRules(currentTable, currentChain)
        self.rulesList = self.formatRules(rawRules)

        stringRules = []
        for rule in self.rulesList :
            stringRules.append('\t'.join(rule))

        self.holdRulesSelection = True
        self.lstRules.clear()
        self.lstRules.addItems(stringRules)
        self.holdRulesSelection = False

        if self.selectFirstRule :
            self.lstRules.setCurrentItem(self.lstRules.item(0))


    @pyqtSlot()
    def on_btnChainsApply_clicked(self):
        if self.lstChains.currentRow() == -1 :
            return

        tableName = str(self.cmbTables.currentText())
        currentChain = self.chainsList[self.lstChains.currentRow()][0]
        policy = str(self.cmbChainsAction.currentText()).upper()

        if self.Firewall.setPolicy(tableName, currentChain, policy) :
            self.refreshChains()
        else :
            QMessageBox.critical(self, "Error!", "Failed to Change the chain's policy.")


    @pyqtSlot()
    def on_btnChainsAdd_clicked(self):
        newChainWindow = NewChain(self)
        newChainWindow.show()


    @pyqtSlot()
    def on_btnChainsDel_clicked(self):
        if self.lstChains.currentRow() == -1 :
            return

        tableName = str(self.cmbTables.currentText())
        currentChain = self.chainsList[self.lstChains.currentRow()][0]

        if self.Firewall.delChain(tableName, currentChain) :
            self.refreshChains()
        else :
            QMessageBox.critical(self, "Error!", "Failed to Delete this chain.")

    @pyqtSlot()
    def on_btnChainsFlush_clicked(self):
        if self.lstChains.currentRow() == -1 :
            return

        tableName = str(self.cmbTables.currentText())
        currentChain = self.chainsList[self.lstChains.currentRow()][0]

        if self.Firewall.flushChain(tableName, currentChain) :
            self.refreshChains()
        else :
            QMessageBox.critical(self, "Error!", "Failed to Flush this chain.")





    #####################################################################
    #                                                                   #
    #                      Handling signals                             #
    #                           Rules                                   #
    #####################################################################


    @pyqtSlot('QString')
    def on_cmbRulesAction_currentIndexChanged(self, selection):
        targetName = str(selection)

        if targetName == 'Run Chain' :
            customeChains = []

            for chain in self.chainsList :
                if chain[1] == 'User defined' :
                    customeChains.append(chain[0])

            self.cmbRulesSpecialAction.clear()
            self.cmbRulesSpecialAction.setEnabled(True)
            self.cmbRulesSpecialAction.addItems(customeChains)

        else :
            self.cmbRulesSpecialAction.clear()
            self.cmbRulesSpecialAction.addItem('Special Action')
            self.cmbRulesSpecialAction.setEnabled(False)

    def on_lstRules_itemSelectionChanged(self):
        if self.holdRulesSelection == True :
            return

        if self.selectedARule == False :
            self.cmbRulesAction.removeItem(0)
            self.selectedARule = True

        ruleID = self.lstRules.currentRow()
        currentAction = self.rulesList[ruleID][0]

        actionsList = ['No action', 'ACCEPT', 'DROP', 'REJECT', \
                        'QUEUE', 'RETURN', 'LOG']

        if actionsList.count(currentAction) == 0 :
            if self.chainsList.count((currentAction, 'User defined')) != 0 :
                self.cmbRulesAction.setCurrentIndex(self.cmbRulesAction.count() - 1)

                customeChains = []
                for chain in self.chainsList :
                    if chain[1] == 'User defined' :
                        customeChains.append(chain[0])

                self.cmbRulesSpecialAction.clear()
                self.cmbRulesSpecialAction.setEnabled(True)
                self.cmbRulesSpecialAction.addItems(customeChains)
                specialRulesIndex = customeChains.index(currentAction)
                self.cmbRulesSpecialAction.setCurrentIndex(specialRulesIndex)

        else :
            ruleIndex = actionsList.index(currentAction)
            self.cmbRulesAction.setCurrentIndex(ruleIndex)

            self.cmbRulesSpecialAction.clear()
            self.cmbRulesSpecialAction.addItem('Special Action')
            self.cmbRulesSpecialAction.setEnabled(False)


    @pyqtSlot()
    def on_btnRulesApply_clicked(self):
        if self.lstRules.currentRow() == -1 :
            return

        tableName = str(self.cmbTables.currentText())
        currentChain = self.chainsList[self.lstChains.currentRow()][0]
        ruleID = str(self.lstRules.currentRow() + 1)
        targetList = ['noTarget', 'ACCEPT', 'DROP', 'REJECT', 'QUEUE', 'RETURN', 'LOG']
        targetID = self.cmbRulesAction.currentIndex()

        if targetID < len(targetList) :
            target = targetList[targetID]

        else :
            target = str(self.cmbRulesSpecialAction.currentText())

        if self.Firewall.moveRule(tableName, currentChain, ruleID, ruleID, target) :
            self.refreshRules('change')

        else :
            QMessageBox.critical(self, "Error!", "Invalid rule target.")


    @pyqtSlot()
    def on_btnRulesAdd_clicked(self):
        ruleDialog = RuleEdit(self)
        ruleDialog.show()

    @pyqtSlot()
    def on_btnRulesDel_clicked(self):
        if self.lstRules.currentRow() == -1 :
            return

        else :
            tableName = str(self.cmbTables.currentText())
            currentChain = self.chainsList[self.lstChains.currentRow()][0]
            ruleID = str(self.lstRules.currentRow() + 1)
            if self.Firewall.deleteRule(tableName, currentChain, ruleID) :
                self.refreshRules('del')

            else :
                QMessageBox.critical(self, "Error!", "Failed to delete this rule.")


    @pyqtSlot()
    def on_btnRulesUp_clicked(self):
        self.moveRule('up')

    @pyqtSlot()
    def on_btnRulesDown_clicked(self):
        self.moveRule('down')


class NewChain(QDialog):

    def __init__(self, parent = None):
        super(NewChain, self).__init__(parent)

        self.parentObject = parent

        self.txtChainName = QLineEdit('Chain Name')
        self.txtChainName.selectAll()
        self.btnApply = QPushButton('Add')
        self.vlMain = QHBoxLayout()

        self.vlMain.addWidget(self.txtChainName)
        self.vlMain.addWidget(self.btnApply)
        self.setLayout(self.vlMain)

        self.setWindowTitle('Add a new chain')

        self.connect(self.btnApply, SIGNAL('clicked()'), self.addChain)

    def addChain(self):
        newChainName = str(self.txtChainName.text())
        currentTable = str(self.parentObject.cmbTables.currentText())


        if self.parentObject.Firewall.addChain(currentTable, newChainName) :
            self.parentObject.refreshChains()
            self.close()

        else :
            self.txtChainName.setText('Invalid Name!')


class RuleEdit(FirewallRuleEditBase.Ui_FirewallRuleEdit, QDialog):
    def __init__(self, parent = None):
        super(RuleEdit, self).__init__(parent)
        self.parentObject = parent
        self.setupUi(self)

        if dir(self.parentObject).count('chainsList') == 0 :
            self.parentObject.on_btnListChains_clicked()

        self.currentChain = self.parentObject.chainsList[\
                            self.parentObject.lstChains.currentRow()][0]

        self.currentOptions = None
        if self.currentChain == 'INPUT' or self.currentChain == 'PREROUTING' :
            self.txtOutputInterface.setEnabled(False)
            self.currentOptions = 'i'

        elif self.currentChain == 'OUTPUT' or self.currentChain == 'POSTROUTING' :
            self.txtInputInterface.setEnabled(False)
            self.currentOptions = 'o'

        elif self.currentChain == 'FORWARD' :
            self.currentOptions = 'io'


    def formatAddress(self, address):
        if address == '' or address == '0.0.0.0/0' :
            return '0.0.0.0/0'

        ipAddr = None
        netmask = None

        if address.count('/') :
            ipList = address.split('/')
            if ipList[0] :
                ipAddr = ipList[0]
            else :
                return False

            if ipList[1] :
                netmask = int(ipList[1])
            else :
                return False

        else :
            ipAddr = address

        if netmask :
            if netmask > 32 or netmask < 0 :
                return False

        ipPortions = ipAddr.split('.')
        if len(ipPortions) != 4 :
            return False

        for portion in ipPortions :
            if int(portion) > 255 or int(portion) < 0 :
                return False

        return address


    def formatInterfaceName(self, ifaceName):
        if ifaceName == '' or ifaceName == 'any' or ifaceName == '*' :
            return 'any'
        else :
            return ifaceName

    def fillCmbSpecialAction(self):
        if dir(self.parentObject).count('chainsList') == 0 :
            self.parentObject.on_btnListChains_clicked()

        self.cmbRulesSpecialAction.clear()
        self.cmbRulesSpecialAction.setEnabled(True)

        customeChains = []
        for chain in self.parentObject.chainsList :
            if chain[1] == 'User defined' :
                customeChains.append(chain[0])

        self.cmbRulesSpecialAction.addItems(customeChains)

    @pyqtSlot('int')
    def on_cmbRulesAction_currentIndexChanged(self, actionID):
        if actionID < (self.cmbRulesAction.count() - 1) :
            self.cmbRulesSpecialAction.clear()
            self.cmbRulesSpecialAction.addItem('Special Action')
            self.cmbRulesSpecialAction.setEnabled(False)

        else :
            self.cmbRulesSpecialAction.setEnabled(True)
            self.fillCmbSpecialAction()


    def on_btnbxApply_accepted(self):
        sourceAddress = str(self.txtSourceAddress.text())
        destinationAddress = str(self.txtDestinationAddress.text())

        sourceAddress = self.formatAddress(sourceAddress)
        if not sourceAddress :
            self.lblNotifications.setText('Bad Source Address!')
            return

        destinationAddress = self.formatAddress(destinationAddress)
        if not destinationAddress :
            self.lblNotifications.setText('Bad Destination Address!')
            return


        inputInterface = str(self.txtInputInterface.text())
        outputInterface = str(self.txtOutputInterface.text())

        inputInterface = self.formatInterfaceName(inputInterface)
        outputInterface = self.formatInterfaceName(outputInterface)

        targetList = ['ACCEPT', 'DROP', 'REJECT', \
                        'QUEUE', 'RETURN', 'LOG']
        target = None
        targetIndex = self.cmbRulesAction.currentIndex()

        if targetIndex > 0 and targetIndex < 7 :
            target = targetList[targetIndex - 1]

        elif targetIndex == 7 :
            target = str(self.cmbRulesSpecialAction.currentText())


        protocol = str(self.cmbProtocol.currentText())
        currentTable = str(self.parentObject.cmbTables.currentText())

        if self.currentOptions == 'i' :
            interfaceOptions = ' -i ' + inputInterface

        elif self.currentOptions == 'o' :
            interfaceOptions = ' -o ' + outputInterface

        elif self.currentOptions == 'io' :
            interfaceOptions = ' -i ' + inputInterface + ' -o ' + outputInterface

        else :
            interfaceOptions = ' -i ' + inputInterface + ' -o ' + outputInterface

        if self.parentObject.Firewall.appendRule(currentTable, self.currentChain, \
                            ' -s ' + sourceAddress + ' -d ' + destinationAddress\
                            + interfaceOptions + ' -p ' + protocol, target) :

            self.parentObject.refreshRules('add')
            self.close()

        else :
            self.lblNotifications.setText('Failed to Add this Rule')







